package com.xf.day05
import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

object TestCombineByKey2 {
  def main(args: Array[String]): Unit = {

    // 创建 SparkConf 对象
    val conf = new SparkConf()
      .setAppName("CombineByKeyWithPartitionLogging")
      .setMaster("local[*]")
      .set("spark.ui.port", "8080")
      .set("spark.driver.host", "127.0.0.1")

    // 创建 SparkContext 对象
    val sc = new SparkContext(conf)
    sc.setLogLevel("ERROR")

    // 创建初始 RDD，3 个分区
    val rdd11 = sc.parallelize(1 to 9, 3)

    // 映射为 (key, value)，key = item % 3
    val rdd12 = rdd11.map(item => (item % 3, item))
    val result1: RDD[(Int, Double)] = rdd12.mapValues(v => v.toDouble)

    // 查看分区数
    println(s"实际分区数: ${result1.getNumPartitions}")
    println()

    // 打印原始分区数据
    println("=== 原始分区数据分布 ===")
    result1.mapPartitionsWithIndex { (pid, iter) =>
      val data = iter.toList
      println(s"分区 $pid: ${data.mkString("[", ", ", "]")}")
      data.iterator
    }.count()
    println()

    // ======== 步骤1：在每个分区内进行局部聚合，并打印带分区号的日志 ========
    // 我们将数据变为 (key, (sum, count), 分区ID)，用于观察过程
    val partialAggregated: RDD[(Int, (Double, Int))] = result1.mapPartitionsWithIndex { (partitionIndex, iterator) =>
      // 使用 mutable Map 模拟 combineByKey 的分区内聚合
      import scala.collection.mutable

      val combinerMap = mutable.Map[Int, (Double, Int)]()

      println(s"--- 开始处理 分区 $partitionIndex (分区内聚合) ---")

      iterator.foreach { case (key, value) =>
        val current = combinerMap.get(key)
        val newAcc = current match {
          case None =>
            println(s"  分区 $partitionIndex | key=$key | 创建初始累加器: value=$value -> (sum=$value, count=1)")
            (value, 1)

          case Some((sum, count)) =>
            val newSum = sum + value
            val newCount = count + 1
            println(s"  分区 $partitionIndex | key=$key | 合并值: 当前(sum=$sum, count=$count) + value=$value -> (sum=$newSum, count=$newCount)")
            (newSum, newCount)
        }
        combinerMap(key) = newAcc
      }

      println(s"--- 分区 $partitionIndex 聚合结果: ${combinerMap.mkString(", ")} ---")
      println()

      // 返回该分区的中间结果
      combinerMap.iterator
    }

    // ======== 步骤2：定义分区间合并函数（用于 combineByKey 的 reduce 阶段）========
    def mergeCombiners(c1: (Double, Int), c2: (Double, Int)): (Double, Int) = {
      val newSum = c1._1 + c2._1
      val newCount = c1._2 + c2._2
      // 这里无法知道来自哪个分区，但可以打印合并行为
      println(s"  [跨分区合并] (sum=${c1._1}, count=${c1._2}) + (sum=${c2._1}, count=${c2._2}) -> (sum=$newSum, count=$newCount)")
      (newSum, newCount)
    }

    // ======== 步骤3：使用 combineByKey 完成分区间合并 ========
    // 注意：partialAggregated 已经是每个分区的聚合结果
    // 现在我们用 combineByKey 的 reduce 阶段合并这些结果（虽然其实可以直接 reduceByKey）

    println("=== 开始分区间合并（reduce 阶段） ===")

    // 使用 combineByKey，但 createCombiner 和 mergeValue 不会再被调用（因为已经是 (sum, count)）
    // 我们只关心 mergeCombiners 的日志
    val finalResultRDD: RDD[(Int, (Double, Int))] = partialAggregated.combineByKey(
      identity, // 直接使用已有值作为初始累加器
      mergeCombiners,
      mergeCombiners
    )

    // 触发计算
    val finalResult = finalResultRDD.collect()

    println("=== 分区间合并完成 ===")
    println()
    println("最终结果:")
    finalResult.foreach { case (key, (sum, count)) =>
      println(s"  键 $key: 总和 = $sum, 计数 = $count, 平均值 = ${sum / count}")
    }

    sc.stop()
  }

}
